import ray
from typing import Optional, List
import time

from pydantic import BaseModel, Field

class Parameter(BaseModel):
    d: int
    e: float

class Result(BaseModel):
    r: float

@ray.remote
def ray_task_test(a, b):
    print("ray_task_test other ray")
    return a + b

@ray.remote
def do_some_work(x):
    time.sleep(1) # Replace this is with work you need to do.
    print(f"Doing work with value {x} and sleeping for 1 second...")
    return x

def rayinray(param: Parameter) -> Result:
    if not ray.is_initialized():
        print("ray has init")
    
    # Launch multiple tasks in parallel
    ray_task1 = ray_task_test.remote(param.d, param.e)
    ray_task2 = do_some_work.remote(1)
    ray_task3 = do_some_work.remote(2)
    ray_task4 = do_some_work.remote(3)

    # Get results from all tasks
    result1  = ray.get([ray_task1, ray_task2, ray_task3, ray_task4])
    
    
    return Result(r=result1[0])
